from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
import os
import sys
import socket
from werkzeug.utils import secure_filename
from services.detection_service import DetectionService
from services.masking_service import MaskingService
from services.alert_service import AlertService
from services.log_service import LogService
from config import Config
import base64
from io import BytesIO
from PIL import Image

app = Flask(__name__, static_folder='../frontend', template_folder='../frontend')
CORS(app)

# 配置文件上传
app.config['UPLOAD_FOLDER'] = Config.UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = Config.MAX_CONTENT_LENGTH

# 初始化服务（延迟加载，避免启动失败）
detection_service = None
masking_service = None
alert_service = None
log_service = None

def init_services():
    """延迟初始化服务"""
    global detection_service, masking_service, alert_service, log_service
    if detection_service is None:
        try:
            detection_service = DetectionService()
            masking_service = MaskingService()
            alert_service = AlertService()
            log_service = LogService(Config.LOG_FILE)
            print("✅ 服务初始化成功")
        except Exception as e:
            print(f"⚠️  服务初始化警告: {e}")
            print("⚠️  部分功能可能不可用，但服务仍会启动")
            # 至少初始化基本服务
            masking_service = MaskingService()
            alert_service = AlertService()
            log_service = LogService(Config.LOG_FILE)
            detection_service = None

# 初始化服务
init_services()

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in Config.ALLOWED_EXTENSIONS

@app.route('/')
def index():
    return send_from_directory('../frontend', 'index.html')

@app.route('/api/detect', methods=['POST'])
def detect_sensitive_data():
    """多模态敏感信息检测API"""
    try:
        if detection_service is None:
            return jsonify({'error': '检测服务未初始化，请检查依赖是否安装完整'}), 503
        
        data = request.get_json()
        content_type = data.get('type')  # 'text' 或 'image'
        content = data.get('content')
        
        if not content:
            return jsonify({'error': '内容不能为空'}), 400
        
        # 调用检测服务
        if content_type == 'text':
            result = detection_service.detect_text(content)
        elif content_type == 'image':
            # 解码base64图片
            try:
                if ',' in content:
                    image_data = base64.b64decode(content.split(',')[1])
                else:
                    image_data = base64.b64decode(content)
                image = Image.open(BytesIO(image_data))
                result = detection_service.detect_image(image)
            except Exception as e:
                return jsonify({'error': f'图片解码失败: {str(e)}'}), 400
        else:
            return jsonify({'error': '不支持的内容类型'}), 400
        
        # 记录检测结果
        if log_service:
            log_service.log_detection(result)
        
        # 判断是否需要告警
        if alert_service and result.get('risk_level') == 'high':
            alert_service.send_alert(result)
        
        return jsonify(result), 200
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/mask', methods=['POST'])
def mask_sensitive_data():
    """敏感信息打码处理API"""
    try:
        if masking_service is None:
            return jsonify({'error': '打码服务未初始化'}), 503
        
        data = request.get_json()
        original_content = data.get('content')
        detected_items = data.get('detected_items', [])
        content_type = data.get('type')
        
        if content_type == 'text':
            masked_content = masking_service.mask_text(original_content, detected_items)
        elif content_type == 'image':
            # 处理图片打码
            try:
                if ',' in original_content:
                    image_data = base64.b64decode(original_content.split(',')[1])
                else:
                    image_data = base64.b64decode(original_content)
                image = Image.open(BytesIO(image_data))
                masked_image = masking_service.mask_image(image, detected_items)
                
                # 转换为base64返回
                buffered = BytesIO()
                masked_image.save(buffered, format="PNG")
                masked_content = base64.b64encode(buffered.getvalue()).decode()
                masked_content = f"data:image/png;base64,{masked_content}"
            except Exception as e:
                return jsonify({'error': f'图片处理失败: {str(e)}'}), 400
        else:
            return jsonify({'error': '不支持的内容类型'}), 400
        
        return jsonify({
            'masked_content': masked_content,
            'masked_items_count': len(detected_items)
        }), 200
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/upload', methods=['POST'])
def upload_file():
    """文件上传接口"""
    if 'file' not in request.files:
        return jsonify({'error': '没有上传文件'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': '文件名为空'}), 400
    
    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
        file.save(filepath)
        
        # 检测文件类型并处理
        if detection_service is None:
            return jsonify({'error': '检测服务未初始化，请检查依赖是否安装完整'}), 503
        
        file_ext = filename.rsplit('.', 1)[1].lower()
        if file_ext in ['png', 'jpg', 'jpeg', 'gif']:
            image = Image.open(filepath)
            result = detection_service.detect_image(image)
        else:
            with open(filepath, 'r', encoding='utf-8') as f:
                text_content = f.read()
            result = detection_service.detect_text(text_content)
        
        if log_service:
            log_service.log_detection(result)
        
        if alert_service and result.get('risk_level') == 'high':
            alert_service.send_alert(result)
        
        return jsonify(result), 200
    
    return jsonify({'error': '不支持的文件类型'}), 400

@app.route('/api/logs', methods=['GET'])
def get_logs():
    """获取检测日志"""
    if log_service is None:
        return jsonify({'error': '日志服务未初始化', 'logs': []}), 503
    logs = log_service.get_recent_logs()
    return jsonify({'logs': logs}), 200

if __name__ == '__main__':
    # 获取端口号（默认18081，如果被占用则尝试其他端口）
    default_port = int(os.getenv('PORT', 18081))
    port = default_port
    
    # 检查端口是否被占用
    def is_port_available(port):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            try:
                s.bind(('', port))
                return True
            except OSError:
                return False
    
    # 如果默认端口被占用，尝试其他端口
    if not is_port_available(port):
        print(f"⚠️  端口 {port} 被占用，尝试其他端口...")
        for p in range(18082, 18091):
            if is_port_available(p):
                port = p
                break
        if port == default_port:
            print(f"❌ 无法找到可用端口，请手动指定端口: PORT=8080 python app.py")
            sys.exit(1)
        else:
            print(f"✅ 使用端口: {port}")
    
    print("🚀 启动多模态敏感数据检测系统...")
    print(f"📝 访问 http://localhost:{port} 使用系统")
    print("按 Ctrl+C 停止服务")
    app.run(debug=True, host='0.0.0.0', port=port)

